package com.ask.forMe.langchain4j.rag.splitter;

import com.ask.forMe.langchain4j.model.vo.SplitterVO;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.splitter.DocumentByCharacterSplitter;
import dev.langchain4j.data.document.splitter.DocumentByParagraphSplitter;
import dev.langchain4j.data.document.splitter.DocumentByRegexSplitter;
import dev.langchain4j.data.document.splitter.DocumentBySentenceSplitter;
import dev.langchain4j.data.segment.TextSegment;

import java.util.List;

/**
 * langchain4j的分割器工具类
 */
public class LangchainSplitter {
    /**
     * 段落分割器
     *
     * @param text
     * @param maxSegmentSize
     * @param maxOverlapSize
     * @return
     */
    public static List<TextSegment> splitByParagraph(String text, int maxSegmentSize, int maxOverlapSize) {
        DocumentByParagraphSplitter splitter = new DocumentByParagraphSplitter(maxSegmentSize, maxOverlapSize);
        List<TextSegment> segments = splitter.split(Document.document(text));
        return segments;
    }

    /**
     * 句子分割器
     *
     * @param text
     * @param maxSegmentSize
     * @param maxOverlapSize
     * @return
     */
    public static List<TextSegment> splitBySentence(String text, int maxSegmentSize, int maxOverlapSize) {
        DocumentBySentenceSplitter splitter = new DocumentBySentenceSplitter(maxSegmentSize, maxOverlapSize);
        List<TextSegment> segments = splitter.split(Document.document(text));
        return segments;
    }

    /**
     * 字符分割器
     *
     * @param text
     * @param maxSegmentSize
     * @param maxOverlapSize
     * @return
     */
    public static List<TextSegment> splitByCharacter(String text, int maxSegmentSize, int maxOverlapSize) {
        DocumentByCharacterSplitter splitter = new DocumentByCharacterSplitter(maxSegmentSize, maxOverlapSize);
        List<TextSegment> segments = splitter.split(Document.document(text));
        return segments;
    }

    /**
     * 自定义分割方式by Regex
     *
     * @param text
     * @param regex
     * @param joinDelimiter
     * @param maxSegmentSize
     * @param maxOverlapSize
     * @return
     */
    public static List<TextSegment> splitByRegex(String text, String regex, String joinDelimiter, int maxSegmentSize, int maxOverlapSize) {
        DocumentByRegexSplitter splitter =
                new DocumentByRegexSplitter(regex, joinDelimiter, maxSegmentSize, maxOverlapSize);
        List<TextSegment> segments = splitter.split(Document.document(text));
        return segments;
    }

    /**
     * 根据type选择不同的分割器1
     *
     * @param type
     * @param text
     * @param splitVO
     * @return
     */
    public static List<TextSegment> splitterSelector(String type, String text, SplitterVO splitVO) {
        if (type.equals("para"))
            return splitByParagraph(text, splitVO.getMaxSegmentSizeInChars(), splitVO.getMaxOverlapSizeInChars());
        else if (type.equals("regex"))
            return splitByRegex(text, splitVO.getRegex(), splitVO.getDelimiter(), splitVO.getMaxSegmentSizeInChars(), splitVO.getMaxOverlapSizeInChars());
        else if (type.equals("sentence"))
            return splitBySentence(text, splitVO.getMaxSegmentSizeInChars(), splitVO.getMaxOverlapSizeInChars());
        else if (type.equals("char"))
            return splitByCharacter(text, splitVO.getMaxSegmentSizeInChars(), splitVO.getMaxOverlapSizeInChars());
        return null;
    }

}
